Metric refactor#524
Conversation
… package QuantityMatrix is renamed to QMatrix for brevity. The monolithic quantity_matrix.py module is supplemented by a new quantity_matrix/ package that splits the implementation into focused submodules (_quantity_matrix, _units_matrix, _det, _inv, _register_primitives, _utils). det and inv operations are added to the public surface. The .py file and the package coexist; Python resolves to the package at import time.
…rix types, and dispatch API Adds a new coordinax._src.metric package with four modules: field.py (AbstractMetricField, AbstractDiagonalMetricField, RoundMetric), matrix.py (AbstractMetricMatrix, DiagonalMetric, DenseMetric), api.py (metric_matrix, metric_representation fallback dispatches), and __init__.py. AbstractMetric/AbstractDiagonalMetric in base/metric.py are renamed to AbstractMetricField/AbstractDiagonalMetricField and stripped of their metric_matrix method — that computation is now a standalone dispatch function.
…e_factors, and rename metric classes Each geometry (Euclidean, HyperSpherical, Minkowski, Product, Embedded, Custom) gains a register_metric.py module that registers metric_matrix and metric_representation dispatches using the new standalone dispatch API. Metric class names are updated: EuclideanMetric → FlatMetric, HyperSphericalMetric → RoundMetric, CartesianProductMetric → ProductMetric, InducedMetric → PullbackMetric. Spherical and Minkowski geometries receive new scale_factors.py files. The metric_matrix method is removed from each metric class in favour of the dispatch function.
…ic symbols The coordinax.api.manifolds dispatch API gains metric_matrix and metric_representation. The coordinax.manifolds re-export module is updated to expose the renamed metric classes (FlatMetric, RoundMetric, ProductMetric, PullbackMetric) and new matrix types (AbstractMetricMatrix, DiagonalMetric, DenseMetric). AbstractManifold is updated to use AbstractMetricField. The scale_factors dispatch file gains metric-aware overloads. Public internal.py and main.py re-exports are aligned with the new names.
…, representations, transforms, vectors, distances, and frames All references to QuantityMatrix (now QMatrix), EuclideanMetric (now FlatMetric), HyperSphericalMetric (now RoundMetric), CartesianProductMetric (now ProductMetric), InducedMetric (now PullbackMetric), AbstractMetric (now AbstractMetricField), and AbstractDiagonalMetric (now AbstractDiagonalMetricField) are updated throughout the core modules and their tests. Docstring Examples section headers that merely introduce the code block are also removed for cleaner inline doctests.
…ng cleanup Remove Examples section headers from docstrings in distance_modulus, frame_transforms, parallax, register_constructors, and register_converters. Update tests to use new metric and frame API conventions.
…ld API Remove unused register_frames dispatches and update tests to match the current frame and manifold API conventions.
…anifold API Update strategy tests and chart/vector/distance tests to reflect the current API conventions introduced by the metric refactor.
…ocstring cleanup Replace QuantityMatrix with QMatrix in qmatrix.py and remove Examples section headers from angles, distances, frames, ptmap, vec_constructors, and vec_converters. Tests updated to match.
spec.md is updated to reflect the new metric module design (AbstractMetricField, metric_matrix dispatch, DiagonalMetric/DenseMetric matrix types). API and guide pages for manifolds, charts, and representations are updated to use the new class names. README and tutorials are adjusted for minor API changes.
Add jax.extend.core and jax.tree_util import abbreviations to pyproject.toml. Add ARG002 (unused method argument) to per-file test ignores to allow pytest fixture parameters that are injected but not referenced in the body.
Cast u.ustrip() results to Array/jax.Array so .ndim/.shape resolve; cast q.unit to u.AbstractUnit; fix ty: ignore placement for J.T attribute access; add ty: ignore[possibly-missing-submodule] on jax.core references.
Replace float literals (1.0, 2.0, etc.) with integer literals where values are whole numbers, and normalize rtol=0.0 to rtol=0 in assert_allclose calls.
…ity_matrix tests Replace float literals like 1.0, 2.0 with int equivalents where the values are whole numbers. Preserve floats where needed for JAX grad/jvp tests, dtype-sensitive operations, and mixed-unit conversion tests that produce float results.
There was a problem hiding this comment.
Pull request overview
This PR is a significant refactor of the metric system in coordinax. It renames types for mathematical clarity (e.g., EuclideanMetric → FlatMetric, HyperSphericalMetric → RoundMetric, QuantityMatrix → QMatrix), restructures the metric API around a three-layer architecture (AbstractMetricField + dispatched metric_matrix function + typed AbstractMetricMatrix results), and introduces pre-defined manifold singletons (cxm.R3, cxm.S2). The diff hunks shown are test-file updates that propagate these renames and adjust assertions to the new typed-result API.
Changes:
- Renames
QuantityMatrix→QMatrix,EuclideanMetric→FlatMetric,HyperSphericalMetric→RoundMetric,InducedMetric→PullbackMetricacross tests. - Updates the scale-factors test to use the new
metric_matrixdispatch andDiagonalMetricresult type instead of the removedmetric.metric_matrix(...)method. - Replaces
cxm.HyperSphericalManifold()/cxm.EuclideanManifold(3)usages with singletonscxm.S2/cxm.R3, and changes various float literals (1.0) to int literals (1) in test data.
Reviewed changes
Copilot reviewed 174 out of 174 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/usage/charts/test_jacobian.py | Renames QuantityMatrix → QMatrix in imports/asserts; changes vmap test inputs to int literals. |
| tests/unit/vectors/test_tangent.py | Switches batched tangent component data to int literals. |
| tests/unit/transforms/conftest.py | Renames QuantityMatrix → QMatrix in fixture; uses int literals for fixture data and expected tuples. |
| tests/unit/representations/test_tangent_map.py | Reformats dicts and converts test data to int literals. |
| tests/unit/manifolds/test_scale_factors_dispatch.py | Adopts new metric_matrix dispatch + DiagonalMetric, renames metric/manifold classes, and renames two test functions to embed QMatrix in their names. |
| tests/unit/manifolds/test_custom.py | Replaces EuclideanMetric with FlatMetric in CustomManifold construction. |
| tests/unit/manifolds/test_angle_between_dispatch.py | Replaces EuclideanMetric/HyperSphericalMetric with FlatMetric/RoundMetric. |
| tests/unit/distances/test_distance.py | Converts float literals to int literals in hypothesis strategies, parametrize cases, and assertions. |
| tests/unit/charts/test_utils.py | Converts float literals in test arrays to int literals. |
…egies and ptmap Replace float literals like 8.0, -8.0, -3.0, 3.0, -5.0, 5.0, 10.0, -10.0 with int equivalents where the values are whole numbers and dtype safety permits. Signed-off-by: nstarman <nstarman@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #524 +/- ##
==========================================
- Coverage 93.24% 92.91% -0.33%
==========================================
Files 232 249 +17
Lines 7386 7877 +491
==========================================
+ Hits 6887 7319 +432
- Misses 499 558 +59 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Metric Refactor Summary
What changed
1. Renamed types for mathematical clarity
EuclideanMetricFlatMetricHyperSphericalMetricRoundMetricInducedMetricPullbackMetricCartesianProductMetricProductMetricAbstractMetricAbstractMetricFieldAbstractDiagonalMetricAbstractDiagonalMetricFieldQuantityMatrixQMatrix2. Three-layer metric architecture replaces method-on-object
Before:
After:
The three layers are:
AbstractMetricField— describes the kind of geometry (flat, round, Lorentzian). Does not compute anything.metric_matrix(manifold, point, chart)— dispatch function that computes and returns the matrix, registered per(manifold, chart)pair.AbstractMetricMatrix(DiagonalMetric/DenseMetric) — typed result carrying the matrix and its sparsity structure.A companion function
metric_representation(manifold, chart)queries which type will be returned without allocating.3. Typed metric matrix results
The old implementation returned a raw
QuantityMatrix(always dense). The new implementation returns:DiagonalMetric— stores only the 1-D diagonal; O(n) matmul instead of O(n²)DenseMetric— full matrix; used for non-orthogonal charts and pullback metricschange_basisautomatically selects the efficient path based on the result type.4.
RoundMetricbecomes an Equinox moduleUnlike the other metrics (frozen static dataclasses),
RoundMetriccan carry a dynamicradiusleaf, enabling JIT andgradthrough radius parameters — important for astronomy use cases.5. Functional-only API for
scale_factorsandangle_betweenPreviously these were also available as convenience methods on metric and manifold objects
(
metric.scale_factors(...),M.angle_between(...)). They are now purely functional, whichsimplifies the class hierarchy and removes the need to call through the object.
6. Pre-defined manifold singletons
Why it's better
metric_matrixmethodQuantityMatrixDiagonalMetricfor orthogonal charts — O(n) ops vs O(n²)metric_matrixdispatch rule without touching existing typesmetric_representation(M, chart)answers without computingEuclideanMetricimplied the manifold;InducedMetrichid the operationFlatMetric,RoundMetric,PullbackMetricuse standard math vocabularyRoundMetriccan carry a differentiable radius